///
/// @file LivenessCFGAnalysis.cpp
/// @brief 活跃变量分析器，基于控制流图的活跃变量分析
/// @author zenglj (zenglj@live.com)
/// @version 1.0
/// @date 2024-11-23
///
/// @copyright Copyright (c) 2024
///
/// @par 修改日志:
/// <table>
/// <tr><th>Date       <th>Version <th>Author  <th>Description
/// <tr><td>2024-11-23 <td>1.0     <td>zenglj  <td>新建
/// </table>
///
#include "Analysis/LivenessCFGAnalysis.h"
#include "Instruction.h"
#include "Instructions/LabelInstruction.h"
#include "Instructions/GotoInstruction.h"
#include "Instructions/CondGotoInstruction.h"
#include <iostream>
#include <algorithm>
#include <unordered_map>

void buildBasicBlocks(const std::vector<Instruction *> & insts, std::vector<BasicBlock *> & blocks)
{
    std::unordered_map<std::string, BasicBlock *> label2block;
    std::unordered_map<Instruction *, BasicBlock *> inst2block;
    BasicBlock * currBlock = nullptr;

    // 第一次遍历：分块，遇到Label新建块
    for (auto * inst: insts) {
        if (inst->getOp() == IRInstOperator::IRINST_OP_LABEL || currBlock == nullptr) {
            currBlock = new BasicBlock();
            blocks.push_back(currBlock);
            if (inst->getOp() == IRInstOperator::IRINST_OP_LABEL) {
                label2block[inst->getIRName()] = currBlock;
            }
        }
        currBlock->addInstruction(inst);
        inst2block[inst] = currBlock;
        // 若遇到终结指令（Goto/CondGoto/Exit），下一个必新块
        if (inst->isTerminator()) {
            currBlock = nullptr;
        }
    }

    // 第二次遍历：建立后继关系
    for (auto * bb: blocks) {
        auto & insts = bb->getInstructions();
        if (insts.empty())
            continue;
        Instruction * last = insts.back();
        if (last->getOp() == IRInstOperator::IRINST_OP_GOTO) {
            auto * gotoInst = static_cast<GotoInstruction *>(last);
            auto * target = gotoInst->getTarget();
            if (target) {
                auto it = label2block.find(target->getIRName());
                if (it != label2block.end()) {
                    bb->addSuccessor(it->second);
                }
            }
        } else if (last->getOp() == IRInstOperator::IRINST_OP_BRANCH) {
            auto * condInst = static_cast<CondGotoInstruction *>(last);
            auto * tLabel = condInst->getTrueLabel();
            auto * fLabel = condInst->getFalseLabel();
            if (tLabel) {
                auto it = label2block.find(tLabel->getIRName());
                if (it != label2block.end()) {
                    bb->addSuccessor(it->second);
                }
            }
            if (fLabel) {
                auto it = label2block.find(fLabel->getIRName());
                if (it != label2block.end()) {
                    bb->addSuccessor(it->second);
                }
            }
        } else if (!last->isTerminator()) {
            // 顺序流：下一个块
            auto it = std::find(blocks.begin(), blocks.end(), bb);
            if (it != blocks.end() && (it + 1) != blocks.end()) {
                bb->addSuccessor(*(it + 1));
            }
        }
    }
}

void LivenessCFGAnalysis::analyze(const std::vector<BasicBlock *> & basicBlocks)
{
    std::map<Instruction *, std::set<std::string>> useSets, defSets;
    for (auto * bb: basicBlocks) {
        for (auto * inst: bb->getInstructions()) {
            std::set<std::string> use, def;
            for (const auto & v: inst->getUseVars())
                use.insert(v);
            for (const auto & v: inst->getDefVars())
                def.insert(v);
            useSets[inst] = use;
            defSets[inst] = def;
        }
    }

    inSets.clear();
    outSets.clear();
    for (auto * bb: basicBlocks) {
        for (auto * inst: bb->getInstructions()) {
            inSets[inst] = {};
            outSets[inst] = {};
        }
    }

    bool changed;
    do {
        changed = false;
        for (auto itBB = basicBlocks.rbegin(); itBB != basicBlocks.rend(); ++itBB) {
            auto * bb = *itBB;
            auto & insts = bb->getInstructions();
            for (int i = (int) insts.size() - 1; i >= 0; --i) {
                Instruction * inst = insts[i];
                std::set<std::string> old_in = inSets[inst];
                std::set<std::string> old_out = outSets[inst];

                std::set<std::string> out;
                if (i + 1 < (int) insts.size()) {
                    out = inSets[insts[i + 1]];
                } else {
                    for (auto * succ: bb->getSuccessors()) {
                        if (!succ->getInstructions().empty()) {
                            Instruction * first = succ->getInstructions().front();
                            out.insert(inSets[first].begin(), inSets[first].end());
                        }
                    }
                }
                outSets[inst] = out;

                std::set<std::string> in = useSets[inst];
                for (const auto & v: out) {
                    if (defSets[inst].count(v) == 0) {
                        in.insert(v);
                    }
                }
                inSets[inst] = in;

                if (in != old_in || out != old_out) {
                    changed = true;
                }
            }
        }
    } while (changed);
}

const std::set<std::string> & LivenessCFGAnalysis::getInSet(Instruction * inst) const
{
    static std::set<std::string> empty;
    auto it = inSets.find(inst);
    return it != inSets.end() ? it->second : empty;
}

void LivenessCFGAnalysis::printLiveness(const std::vector<BasicBlock *> & basicBlocks) const
{
    int blockIdx = 0;
    for (auto * bb: basicBlocks) {
        // 块首指令
        Instruction * first = bb->getInstructions().empty() ? nullptr : bb->getInstructions().front();
        Instruction * last = bb->getInstructions().empty() ? nullptr : bb->getInstructions().back();

        // 计算use[B]和def[B]
        std::set<std::string> useB, defB;
        for (auto * inst: bb->getInstructions()) {
            for (const auto & v: inst->getUseVars()) {
                if (defB.count(v) == 0)
                    useB.insert(v);
            }
            for (const auto & v: inst->getDefVars()) {
                defB.insert(v);
            }
        }

        // 块入口in[B]
        std::cout << "---- BasicBlock #" << blockIdx++;
        if (first && first->getOp() == IRInstOperator::IRINST_OP_LABEL) {
            std::cout << " (Label: " << first->getIRName() << ")";
        }
        std::cout << " ----" << std::endl;
        if (first) {
            std::cout << "in[B]: { ";
            for (const auto & v: getInSet(first))
                std::cout << v << " ";
            std::cout << "}" << std::endl;
        }
        std::cout << "use[B]: { ";
        for (const auto & v: useB)
            std::cout << v << " ";
        std::cout << "}" << std::endl;
        std::cout << "def[B]: { ";
        for (const auto & v: defB)
            std::cout << v << " ";
        std::cout << "}" << std::endl;

        // 每条指令in集合
        for (auto * inst: bb->getInstructions()) {
            std::string ir;
            inst->toString(ir);
            std::cout << "[LIVE IN] { ";
            for (const auto & v: getInSet(inst)) {
                std::cout << v << " ";
            }
            std::cout << "}  |  " << ir << std::endl;
        }

        // 块出口out[B]
        if (last) {
            std::cout << "out[B]: { ";
            auto it = outSets.find(last);
            if (it != outSets.end()) {
                for (const auto & v: it->second)
                    std::cout << v << " ";
            }
            std::cout << "}" << std::endl;
        }
    }
}